Home Explore Help
Register Sign In
J1X
/
huginn
Watch 2
Star 0
Fork 0

Browse Source
use TwitterConcern in the twitter_stream_agent; cleanup unused memory; add specs

Andrew Cantino 11 years ago
parent
8d864c4a0b
commit
0dcee80cba
4 changed files with 127 additions and 25 deletions
Show Diff Stats
  1. 2 3
      app/concerns/twitter_concern.rb
  2. 25 18
      app/models/agents/twitter_stream_agent.rb
  3. 4 4
      bin/twitter_stream.rb
  4. 96 0
      spec/models/agents/twitter_stream_agent_spec.rb

+ 2 - 3
app/concerns/twitter_concern.rb
View File 

@@ -19,12 +19,11 @@ module TwitterConcern
19 19 
     Twitter.configure do |config|
20 20 
       config.consumer_key = options[:consumer_key]
21 21 
       config.consumer_secret = options[:consumer_secret]
22 
-      config.oauth_token = options[:oauth_token]
23 
-      config.oauth_token_secret = options[:oauth_token_secret]
22 
+      config.oauth_token = options[:oauth_token] || options[:access_key]
23 
+      config.oauth_token_secret = options[:oauth_token_secret] || options[:access_secret]
24 24 
     end
25 25 
   end
26 26
27 27 
   module ClassMethods
28 
-
29 28 
   end
30 29 
 end

+ 25 - 18
app/models/agents/twitter_stream_agent.rb
View File 

@@ -1,11 +1,12 @@
1 1 
 module Agents
2 2 
   class TwitterStreamAgent < Agent
3 
+    include TwitterConcern
3 4 
     cannot_receive_events!
4 5
5 6 
     description <<-MD
6 7 
       The TwitterStreamAgent follows the Twitter stream in real time, watching for certain keywords, or filters, that you provide.
7 8
8 
-      You must provide an oAuth `consumer_key`, `consumer_secret`, `access_key`, and `access_secret`, as well as an array of `filters`.  Multiple words in a filter
9 
+      You must provide an oAuth `consumer_key`, `consumer_secret`, `oauth_token`, and `oauth_token_secret`, as well as an array of `filters`.  Multiple words in a filter
9 10 
       must all show up in a tweet, but are independent of order.
10 11
11 12 
       To get oAuth credentials for Twitter, [follow these instructions](https://github.com/cantino/huginn/wiki/Getting-a-twitter-oauth-token).
 
@@ -51,14 +52,10 @@ module Agents
51 52 
     default_schedule "11pm"
52 53
53 54 
     def validate_options
54 
-      unless options[:consumer_key].present? &&
55 
-             options[:consumer_secret].present? &&
56 
-             options[:access_key].present? &&
57 
-             options[:access_secret].present? &&
58 
-             options[:filters].present? &&
55 
+      unless options[:filters].present? &&
59 56 
              options[:expected_update_period_in_days].present? &&
60 57 
              options[:generate].present?
61 
-        errors.add(:base, "expected_update_period_in_days, generate, consumer_key, consumer_secret, access_key, access_secret, and filters are required fields")
58 
+        errors.add(:base, "expected_update_period_in_days, generate, and filters are required fields")
62 59 
       end
63 60 
     end
64 61 
 
@@ -70,8 +67,8 @@ module Agents
70 67 
       {
71 68 
           :consumer_key => "---",
72 69 
           :consumer_secret => "---",
73 
-          :access_key => "---",
74 
-          :access_secret => "---",
70 
+          :oauth_token => "---",
71 
+          :oauth_token_secret => "---",
75 72 
           :filters => %w[keyword1 keyword2],
76 73 
           :expected_update_period_in_days => "2",
77 74 
           :generate => "events"
 
@@ -80,24 +77,34 @@ module Agents
80 77
81 78 
     def process_tweet(filter, status)
82 79 
       if options[:generate] == "counts"
83 
-        # Avoid memory pollution
84 
-        me = Agent.find(id)
85 
-        me.memory[:filter_counts] ||= {}
86 
-        me.memory[:filter_counts][filter.to_sym] ||= 0
87 
-        me.memory[:filter_counts][filter.to_sym] += 1
88 
-        me.save!
80 
+        # Avoid memory pollution by reloading the Agent.
81 
+        agent = Agent.find(id)
82 
+        agent.memory[:filter_counts] ||= {}
83 
+        agent.memory[:filter_counts][filter.to_sym] ||= 0
84 
+        agent.memory[:filter_counts][filter.to_sym] += 1
85 
+        remove_unused_keys!(agent, :filter_counts)
86 
+        agent.save!
89 87 
       else
90 88 
         create_event :payload => status.merge(:filter => filter.to_s)
91 89 
       end
92 90 
     end
93 91
94 92 
     def check
95 
-      if memory[:filter_counts] && memory[:filter_counts].length > 0
93 
+      if options[:generate] == "counts" && memory[:filter_counts] && memory[:filter_counts].length > 0
96 94 
         memory[:filter_counts].each do |filter, count|
97 95 
           create_event :payload => { :filter => filter.to_s, :count => count, :time => Time.now.to_i }
98 96 
         end
99 
-        memory[:filter_counts] = {}
100 
-        save!
97 
+      end
98 
+      memory[:filter_counts] = {}
99 
+    end
100 
+
101 
+    protected
102 
+
103 
+    def remove_unused_keys!(agent, base)
104 
+      if agent.memory[base]
105 
+        (agent.memory[base].keys - agent.options[:filters].map(&:to_sym)).each do |removed_key|
106 
+          agent.memory[base].delete(removed_key)
107 
+        end
101 108 
       end
102 109 
     end
103 110 
   end

+ 4 - 4
bin/twitter_stream.rb
View File 

@@ -17,11 +17,11 @@ require 'pp'
17 17 
 def stream!(filters, options = {}, &block)
18 18 
   stream = Twitter::JSONStream.connect(
19 19 
     :path    => "/1/statuses/#{(filters && filters.length > 0) ? 'filter' : 'sample'}.json#{"?track=#{filters.map {|f| CGI::escape(f) }.join(",")}" if filters && filters.length > 0}",
20 
-    :oauth => {
20 
+    :oauth   => {
21 21 
       :consumer_key    => options[:consumer_key],
22 22 
       :consumer_secret => options[:consumer_secret],
23 
-      :access_key      => options[:access_key],
24 
-      :access_secret   => options[:access_secret]
23 
+      :access_key      => options[:oauth_token] || options[:access_key],
24 
+      :access_secret   => options[:oauth_token_secret] || options[:access_secret]
25 25 
     },
26 26 
     :ssl     => true
27 27 
   )
 
@@ -60,7 +60,7 @@ def load_and_run(agents)
60 60 
       end
61 61 
     end
62 62
63 
-    options = agents.first.options.slice(:consumer_key, :consumer_secret, :access_key, :access_secret)
63 
+    options = agents.first.options.slice(:consumer_key, :consumer_secret, :access_key, :oauth_token, :access_secret, :oauth_token_secret)
64 64
65 65 
     recent_tweets = []
66 66

+ 96 - 0
spec/models/agents/twitter_stream_agent_spec.rb
View File 

@@ -0,0 +1,96 @@
1 
+require 'spec_helper'
2 
+
3 
+describe Agents::TwitterStreamAgent do
4 
+  before do
5 
+    @opts = {
6 
+      :consumer_key => "---",
7 
+      :consumer_secret => "---",
8 
+      :oauth_token => "---",
9 
+      :oauth_token_secret => "---",
10 
+      :filters => %w[keyword1 keyword2],
11 
+      :expected_update_period_in_days => "2",
12 
+      :generate => "events"
13 
+    }
14 
+
15 
+    @agent = Agents::TwitterStreamAgent.new(:name => "HuginnBot", :options => @opts)
16 
+    @agent.user = users(:bob)
17 
+    @agent.save!
18 
+  end
19 
+
20 
+  describe '#process_tweet' do
21 
+    context "when generate is set to 'counts'" do
22 
+      before do
23 
+        @agent.options[:generate] = 'counts'
24 
+      end
25 
+
26 
+      it 'records counts' do
27 
+        @agent.process_tweet(:keyword1, {:text => "something", :user => {:name => "Mr. Someone"}})
28 
+        @agent.process_tweet(:keyword2, {:text => "something", :user => {:name => "Mr. Someone"}})
29 
+        @agent.process_tweet(:keyword1, {:text => "something", :user => {:name => "Mr. Someone"}})
30 
+
31 
+        @agent.reload
32 
+        @agent.memory[:filter_counts][:keyword1].should == 2
33 
+        @agent.memory[:filter_counts][:keyword2].should == 1
34 
+      end
35 
+
36 
+      it 'removes unused keys' do
37 
+        @agent.memory[:filter_counts] = {:keyword1 => 2, :keyword2 => 3, :keyword3 => 4}
38 
+        @agent.save!
39 
+        @agent.process_tweet(:keyword1, {:text => "something", :user => {:name => "Mr. Someone"}})
40 
+        @agent.reload.memory[:filter_counts].should == {:keyword1 => 3, :keyword2 => 3}
41 
+      end
42 
+    end
43 
+
44 
+    context "when generate is set to 'events'" do
45 
+      it 'emits events immediately' do
46 
+        lambda {
47 
+          @agent.process_tweet('keyword1', {:text => "something", :user => {:name => "Mr. Someone"}})
48 
+        }.should change { @agent.events.count }.by(1)
49 
+
50 
+        @agent.events.last.payload.should == {
51 
+          :filter => 'keyword1',
52 
+          :text => "something",
53 
+          :user => {:name => "Mr. Someone"}
54 
+        }
55 
+      end
56 
+    end
57 
+  end
58 
+
59 
+  describe '#check' do
60 
+    context "when generate is set to 'counts'" do
61 
+      before do
62 
+        @agent.options[:generate] = 'counts'
63 
+        @agent.save!
64 
+      end
65 
+
66 
+      it 'emits events' do
67 
+        @agent.process_tweet(:keyword1, {:text => "something", :user => {:name => "Mr. Someone"}})
68 
+        @agent.process_tweet(:keyword2, {:text => "something", :user => {:name => "Mr. Someone"}})
69 
+        @agent.process_tweet(:keyword1, {:text => "something", :user => {:name => "Mr. Someone"}})
70 
+
71 
+        lambda {
72 
+          @agent.reload.check
73 
+        }.should change { @agent.events.count }.by(2)
74 
+
75 
+        @agent.events[-1].payload[:filter].should == 'keyword1'
76 
+        @agent.events[-1].payload[:count].should == 2
77 
+
78 
+        @agent.events[-2].payload[:filter].should == 'keyword2'
79 
+        @agent.events[-2].payload[:count].should == 1
80 
+
81 
+        @agent.memory[:filter_counts].should == {}
82 
+      end
83 
+    end
84 
+
85 
+    context "when generate is not set to 'counts'" do
86 
+      it 'does nothing' do
87 
+        @agent.memory[:filter_counts] = { :keyword1 => 2 }
88 
+        @agent.save!
89 
+        lambda {
90 
+          @agent.reload.check
91 
+        }.should_not change { Event.count }
92 
+        @agent.memory[:filter_counts].should == {}
93 
+      end
94 
+    end
95 
+  end
96 
+end